import csv
import re
import pandas as pd
import logging
from typing import List, Tuple

from utils.instances import TOKENIZER, LLM
from utils import prompts
from langchain_core.prompts import ChatPromptTemplate
import utils.configFinRAG as configFinRAG

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def extract_dates(question: str) -> Tuple[str, List[str]]:
    """提取问题中的日期并替换为空格"""
    date_pattern = r'\d{8}'
    dates = re.findall(date_pattern, question)
    question_without_dates = question
    for date in dates:
        question_without_dates = question_without_dates.replace(date, ' ')
    return question_without_dates, dates

def calculate_jaccard_similarity(tokens1: set, tokens2: set) -> float:
    """计算两个token集合的Jaccard相似度"""
    intersection = len(tokens1 & tokens2)
    union = len(tokens1 | tokens2)
    return intersection / union if union > 0 else 0

def find_most_similar_examples(question_tokens: set, example_token_list: List[set], example_num: int) -> List[int]:
    """找到与问题最相似的示例索引"""
    similarity_list = [calculate_jaccard_similarity(question_tokens, tokens) for tokens in example_token_list]
    sorted_indices = sorted(range(len(similarity_list)), key=lambda i: similarity_list[i], reverse=True)
    return sorted_indices[:example_num]

def generate_sql(question: str, llm, example_question_list: List[str], example_sql_list: List[str],
                 example_token_list: List[set], example_num: int = 5) -> Tuple[str, str]:
    """生成SQL查询"""
    sql_pattern_start = '```sql'
    sql_pattern_end = '```'

    # 提取日期并替换为空格
    question_without_dates, _ = extract_dates(question)
    question_tokens = set(TOKENIZER(question_without_dates)['input_ids'])

    # 找到最相似的示例
    similar_indices = find_most_similar_examples(question_tokens, example_token_list, example_num)

    # 组装prompt
    examples = '\n'.join(
        [f"问题：{example_question_list[index]}\nSQL：{example_sql_list[index]}" for index in similar_indices])
    prompt = ChatPromptTemplate.from_template(prompts.GENERATE_SQL_TEMPLATE)
    chain = prompt | llm
    response = chain.invoke({"examples": examples, "table_info": prompts.TABLE_INFO, "question": question})

    # 提取SQL
    sql = response.content
    start_index = sql.find(sql_pattern_start) + len(sql_pattern_start)
    end_index = sql.find(sql_pattern_end, start_index)

    if start_index >= 0 and end_index >= start_index:
        sql = sql[start_index:end_index].strip()
        return prompt.invoke({"examples": examples, "table_info": prompts.TABLE_INFO, "question": question}), sql
    else:
        logging.error(f"生成SQL错误: {question}")
        return "error", "error"

def main():
    # 读取问题和SQL模板
    sql_examples_file = pd.read_csv(configFinRAG.sql_examples_path, delimiter=",", header=0)
    example_question_list = sql_examples_file['问题'].tolist()
    example_sql_list = sql_examples_file['SQL'].tolist()
    example_token_list = [set(TOKENIZER(question)['input_ids']) for question in example_question_list]

    # 读取测试问题
    question_csv_file = pd.read_csv(configFinRAG.question_classify_path, delimiter=",", header=0)

    with open(configFinRAG.question_sql_path, 'w', newline='', encoding='utf-8-sig') as question_sql_file:
        csv_writer = csv.writer(question_sql_file)
        csv_writer.writerow(['问题id', '问题', 'SQL', 'prompt'])

        # 处理查询数据库的问题
        query_questions = question_csv_file[question_csv_file['分类'] == '查询数据库']
        for _, row in query_questions.iterrows():
            result_prompt, result = generate_sql(row['问题'], LLM, example_question_list, example_sql_list,
                                                 example_token_list)
            csv_writer.writerow([str(row['问题id']), str(row['问题']), result, result_prompt])

        # 跳过非查询数据库的问题
        non_query_questions = question_csv_file[question_csv_file['分类'] != '查询数据库']
        for _, row in non_query_questions.iterrows():
            logging.info(f"跳过问题: {row['问题']}")

if __name__ == '__main__':
    main()